package com.gmadmin.utils;

import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author zhangliang
 * @version 1.0
 * @date 2022/6/14 15:32
 */
public class TreeUtil {


    private TreeUtil() {

    }

    /**
     * 树形结构转换并排序
     * Map版本（速度比递归要快很多）
     * listToTree
     *
     * @param target      需转换的数据
     * @param getId       主键
     * @param getParentId 父id （父id必须和主键相同类型）
     * @param getChildren 子集
     * @param setChildren 子集
     * @return tree
     */
    public static <T, R> List<T> listToTree(List<T> target, Function<T, R> getId, Function<T, R> getParentId,
                                            Function<T, List<T>> getChildren, BiConsumer<T, List<T>> setChildren, Function<T, Long> getSort) {

        Map<R, T> oldMap = target.stream().collect(Collectors.toMap(getId, t -> t));
        List<T> result = new ArrayList<>();
        target.forEach(tree -> {
            T parent = oldMap.get(getParentId.apply(tree));
            if (parent == null) {
                result.add(tree);
                Collections.sort(result, (o1, o2) -> (int) (getSort.apply(o1) - getSort.apply(o2)));
            } else {
                List<T> ch = getChildren.apply(parent);
                if (ch == null) {
                    ch = new ArrayList<>();
                }
                ch.add(tree);
                Collections.sort(ch, (o1, o2) -> (int) (getSort.apply(o1) - getSort.apply(o2)));
                setChildren.accept(parent, ch);
            }
        });
        return result;
    }
}
